"""GiW mixture"""

import numpy as np
from giw import Giw
from scipy.special import digamma
from ldmat import *
import pprint


class MixComponent(object):
    """
    pdf mixture component

    Attributes
    ----------
    pdf : Giw
        pdf of the GiW distribution. Instance of Giw
    weight : float
        Weight (probability) of component
    """
    def __init__(self, pdf, weight):
        self.pdf = pdf
        self.weight = weight



class GiwMixture(object):
    """
    Gauss-inverse-Wishart mixtures

    Attributes
    ----------
    _mix: list
        array with components

    Methods
    -------
    __init__()
        constructor
    add(component_pdf, weight)
        Add component (pdf)
    rem(no_of_component)
        Remove component
    approx()
        Approximation of the mixture
    get_weights()
        Returns vector of weghts
    set_weights(weghts)
        Set weights of all components at once
    """

    def __init__(self):
        self._mix = [] 

    def __str__(self):
        rep = "\n" + "=" * 11
        rep += "\n  Mixture\n"
        rep += "=" * 11 + "\n"
        rep += "No. of components: " + str(len(self._mix))
        return rep

    def add(self, component_pdf, weight):
        """Addition of a component into the mixture"""
        self._mix.append(\
            MixComponent(component_pdf, weight)\
        )

    def rem(self, no_of_component):
        """Removal of a component from the mixture"""
        self._mix.remove(component_number)

    def get_weights(self):
        """Returns vector of weghts"""
        return [component.weight for component in self._mix]
    
    def set_weights(self, weights):
        """Set weights of all components at once"""
        for i in xrange(len(self._mix)):
        	self._mix[i].weight = weights[i]

    def approx(self):
        """
        Approximation of the mixture by a single GiW pdf
        Details on var names - see Research report

        Returns
        -------
        Giw instance
        """
        no_of_components = len(self._mix)
        # d_lsr: LS remainders
        d_lsr = np.array([self._mix[i].pdf.v.d[0] \
                for i in xrange(no_of_components)])
        # vector of weights
        weights = np.array([self._mix[i].weight \
                for i in xrange(no_of_components)])
        #vector of nu's 
        nu = np.array([self._mix[i].pdf.nu \
                for i in xrange(no_of_components)])

        # common_part = lambda * nu / D_LSR
        common_part = [self._mix[i].weight \
                       * self._mix[i].pdf.nu \
                       for i in xrange(no_of_components)]
        common_part /= d_lsr
        sum_common_part = np.sum(common_part)

        ################### THETA ##################
        # common_part_theta = sum(common_part * theta)
        common_part_theta = 0
        for i in xrange(no_of_components):
        	common_part_theta += common_part[i] \
        	                     * self._mix[i].pdf.est_theta

        # Approximate theta
        approximate_est_theta = 1/np.sum(common_part) * common_part_theta


        ################### Nu ##################
        # substitution
        a = np.log(sum_common_part) \
            + np.dot(weights, np.log(d_lsr)) \
            - np.dot(weights, digamma(0.5*nu))

        approximate_nu = 1. + np.sqrt(1. + 4./3.*(a - np.log(2.))) 
        approximate_nu /= 2.*(a - np.log(2.))

        
        ################### D_LSR ##################
        approximate_d_lsr = approximate_nu * 1/sum_common_part

        
        ################### Cov(theta) ##################
        # disgusting for cycles should be omitted in the future
        approximate_est_theta_cov = 0
        for i in xrange(no_of_components):
        	theta_difference = self._mix[i].pdf.est_theta - approximate_est_theta
        	approximate_est_theta_cov += weights[i] * self._mix[i].pdf.est_theta_cov
        	approximate_est_theta_cov += common_part[i] \
        	    * np.dot(theta_difference, theta_difference)

        
        ################### Minimizer pdf ###############
        apr_cov_l, apr_cov_d = ltdl(np.linalg.inv(approximate_est_theta_cov))
        approx_d = np.r_[approximate_d_lsr, apr_cov_d]

        approx_l = np.eye(len(approx_d))
        approx_l_tmp = np.dot(apr_cov_l, approximate_est_theta)
        approx_l_tmp = np.column_stack((approx_l_tmp, apr_cov_l))
        approx_l[1:] = approx_l_tmp

        approx_ld = approx_l + np.diag(approx_d) - np.eye(len(approx_d))
        
        return Giw(Ldmat(approx_ld, from_ld=True), \
                   approximate_nu, desc="Mixture approximate")
        


#===========================================
if __name__ == "__main__":
    TestMat = np.array([
                   [ 0.7 ,  0.4 ,  0.  ],
                   [ 0.4 ,  0.5 ,  0.15],
                   [ 0.  ,  0.15,  0.6 ]])
    GiwTest = Giw(TestMat)
    GiwTest2 = Giw(TestMat)

    MixTest = GiwMixture()
    MixTest.add(GiwTest, 0.1)
    MixTest.add(GiwTest2, 0.9)
    print MixTest.get_weights()
    MixTest.set_weights([0.8, 0.2])
    print MixTest.get_weights()

    Apr = MixTest.approx()

    print "======================"
    print "DEBUG"
    print "LD - LDapr\n", np.linalg.det(GiwTest.v.v - Apr.v.v)
    print "nu - nuapr", GiwTest.nu - Apr.nu
    print "est_theta - est_theta_apr", GiwTest.est_theta - Apr.est_theta
    print "est_theta_cov - est_theta_cov_apr\n", np.linalg.det(GiwTest.est_theta_cov - Apr.est_theta_cov)

